import math, os, sys, poser, random

os.chdir( os.path.dirname(poser.AppLocation()) )
sys.path.append(os.path.join("Runtime","Python","Poseworks","Particles3","data"))

from prt3Constants import *
from prt3Math import *
import prt3Core

class PwPoserMaterial(object):
    def __init__(self, obj):
        self.name = obj.Name()
        self.KdColor = obj.DiffuseColor()
        self.KaColor = obj.AmbientColor()
        self.KsColor = obj.SpecularColor()
        self.tMax = obj.TransparencyMax()
        self.textureMap = obj.TextureMapFileName()
        if not self.textureMap:
            self.textureMap = "NO_MAP"
        else:
            self.textureMap = "\"" + self.textureMap + "\""
        self.transparencyMap = obj.TransparencyMapFileName()
        if not self.transparencyMap:
            self.transparencyMap = "NO_MAP"
        else:
            self.transparencyMap = "\"" + self.transparencyMap + "\""
        self.shaderTree = PwPoserShaderTree(obj.ShaderTree())
    def Name(self):
        return self.name
    def write(self):
        stData = ""
        for node in self.shaderTree.Nodes():
            ipData = ""
            for ip in node.Inputs():
                ipfile = ""
                value = ip.Value()
                if isinstance(value, str):
                    ipfile = value
                    value = (1, 1, 1)
                elif isinstance(value, float):
                    value = (value, 0, 0)
                elif not value:
                    value = (0, 0, 0)
                innode = ip.InNode()
                if not innode:
                    innode = "NO_NODE"
                ipData = ipData + "\n\t\t\t\tnodeInput \"%s\"" % ip.InternalName() +\
                         "\n\t\t\t\t\t{" +\
                         "\n\t\t\t\t\tname \"%s\"" % ip.Name() +\
                         "\n\t\t\t\t\tvalue %f %f %f" % value +\
                         "\n\t\t\t\t\tparmR NO_PARM" +\
                         "\n\t\t\t\t\tparmG NO_PARM" +\
                         "\n\t\t\t\t\tparmB NO_PARM" +\
                         "\n\t\t\t\t\tnode %s" % innode +\
                         "\n\t\t\t\t\tfile \"%s\"" % ipfile +\
                         "\n\t\t\t\t\t}"
            stData = stData + "\n\t\t\tnode \"%s\" \"%s\"" % (node.Type(), node.InternalName()) +\
                     "\n\t\t\t\t{" +\
                     "\n\t\t\t\tname \"%s\"" % node.Name() +\
                     "\n\t\t\t\tpos %i %i" % node.Location() +\
                     "\n\t\t\t\tshowPreview %i" % node.PreviewCollapsed() +\
                     "\n\t\t\t\tinputsCollapsed %s" % node.InputsCollapsed() +\
                     ipData +\
                     "\n\t\t\t\t}"
        return "\n\tmaterial %s" % self.Name() +\
               "\n\t\t{" +\
               "\n\t\tKdColor %f %f %f 1" % self.KdColor +\
               "\n\t\tKaColor %f %f %f 1" % self.KaColor +\
               "\n\t\tKsColor %f %f %f 1" % self.KsColor +\
               "\n\t\ttMax %f" % self.tMax +\
               "\n\t\ttextureMap %s" % self.textureMap +\
               "\n\t\ttransparencyMap %s" % self.transparencyMap +\
               "\n\t\tshaderTree" +\
               "\n\t\t\t{" +\
               stData +\
               "\n\t\t\t}" +\
               "\n\t\t}"

class PwPoserShaderTree(object):
    def __init__(self, st):
        self.nodes = []
        for node in st.Nodes():
            self.nodes.append( PwPoserNode(node) )
    def Nodes(self):
        return self.nodes

class PwPoserNode(object):
    def __init__(self, node):
        self.inputs = []
        for ip in node.Inputs():
            self.inputs.append( PwPoserInput(ip) )
        self.inputscollapsed = node.InputsCollapsed()
        self.internalname = node.InternalName()
        self.location = node.Location()
        self.name = node.Name()
        self.previewcollapsed = node.PreviewCollapsed()
        self.type = node.Type()
    def Inputs(self):
        return self.inputs
    def InputsCollapsed(self):
        return self.inputscollapsed
    def InternalName(self):
        return self.internalname
    def Location(self):
        return self.location
    def Name(self):
        return self.name
    def PreviewCollapsed(self):
        return self.previewcollapsed
    def Type(self):
        return self.type

class PwPoserInput(object):
    def __init__(self, ip):
        self.animated = ip.Animated()
        self.canbeanimated = ip.CanBeAnimated()
        self.innode = ip.InNode()
        if self.innode:
            self.innode = self.innode.InternalName()
        self.internalname = ip.InternalName()
        self.itsnode = ip.ItsNode().InternalName()
        self.name = ip.Name()
        self.parameters = ip.Parameters()
        self.type = ip.Type()
        self.value = ip.Value()
    def Animated(self):
        return self.animated
    def CanBeAnimated(self):
        return self.canbeanimated
    def InNode(self):
        return self.innode
    def InternalName(self):
        return self.internalname
    def ItsNode(self):
        return self.itsnode
    def Name(self):
        return self.name
    def Parameters(self):
        return self.parameters
    def Type(self):
        return self.type
    def Value(self):
        return self.value

def collectShader(source):
    return PwPoserShaderTree(source.ShaderTree())

def createFromStored(stored, material):
    equivs = {}
    st = material.ShaderTree()
    for node in st.Nodes():
        if node.Type() == poser.kNodeTypeCodePOSERSURFACE:
            continue
        else:
            node.Delete()
    for node in stored.Nodes():
        newNode = None
        if node.Type() == poser.kNodeTypeCodePOSERSURFACE:
            newNode = st.Node(0)
        else:
            newNode = st.CreateNode(node.Type())
        equivs[ node.InternalName() ] = newNode.InternalName()
        newNode.SetName(node.Name())
        newNode.SetLocation(node.Location()[0], node.Location()[1])
        newNode.SetInputsCollapsed(node.InputsCollapsed())
        newNode.SetPreviewCollapsed(node.PreviewCollapsed())
        for ip in node.Inputs():
            newInput = newNode.InputByInternalName(ip.InternalName())
            #newInput.SetAnimated( ip.Animated() )
            newInput.SetName( ip.Name() )
            kind = ip.Type()
            value = ip.Value()
            if kind == poser.kNodeInputCodeCOLOR or kind == poser.kNodeInputCodeVECTOR:
                r, g, b = value
                newInput.SetColor( r, g, b )
            elif kind == poser.kNodeInputCodeSTRING:
                if value == None:
                    value = ""
                newInput.SetString( value )
            else:
                newInput.SetFloat( value )
    for node in stored.Nodes():
        for ip in node.Inputs():
            inNode = ip.InNode()
            if inNode:
                outNode = st.NodeByInternalName( equivs[inNode] )
                inNode = st.NodeByInternalName( equivs[ip.ItsNode()] )
                outNode.ConnectToInput( inNode.InputByInternalName(ip.InternalName()) )
    root = st.Node(0)
    r, g, b = root.InputByInternalName("Diffuse_Color").Value()
    material.SetDiffuseColor(r, g, b)
    r, g, b = root.InputByInternalName("Ambient_Color").Value()
    material.SetAmbientColor(r, g, b)
    r = root.InputByInternalName("Transparency_Max").Value()
    material.SetTransparencyMax(r)

def copyMaterials(actor, figure):
    source = actor.Materials()[0]
    destinations = []
    for act in figure.Actors():
        if act.IsBodyPart():
            material = act.Material("Preview")
            if material:
                destinations.append( material )
        else:
            continue
    src = collectShader(source)
    for dest in destinations:
       createFromStored(src, dest)
    for act in figure.Actors():
        updatePrt3Nodes(act, copy=1)

def findLifeNode(actor):
    nodes = actor.Materials()[0].ShaderTree().Nodes()
    lifeNode = None
    for node in nodes:
        if node.Name() == "PRT3_Life":
            lifeNode = node
            break
    return lifeNode

def configureLifeNode(actor, lifeNode):
    if lifeNode:
        kind = actor.Parameter(PARM_TYPE).Value()
        life = 0.0
        if kind == PARTICLE_TYPE:
            life = actor.Parameter(PARM_LIFE).Value()
        lifeNode.InputByInternalName("Value_1").SetFloat(life)
        body = actor.ItsFigure().ParentActor()
        lifespan = body.Parameter(PARM_LIFESPAN).Value()
        lifeNode.InputByInternalName("Value_2").SetFloat(lifespan)
        lifeNode.InputByInternalName("Math_Argument").SetFloat(4)

def addLifeNode(actor, kind=GROUP_TYPE):
    lifeNode = findLifeNode(actor)
    if not lifeNode:
        material = actor.Materials()[0]
        st = material.ShaderTree()
        lifeNode = st.CreateNode(poser.kNodeTypeCodeMATH)
        lifeNode.SetName("PRT3_Life")
        lifeNode.SetInputsCollapsed(1)
    configureLifeNode(actor, lifeNode)
    return lifeNode

def updateLifeNode(actor):
    lifeNode = findLifeNode(actor)
    if lifeNode:
        configureLifeNode(actor, lifeNode)

def findRandomNodes(actor):
    material = actor.Materials()[0]
    st = material.ShaderTree()
    randNodes = []
    for node in st.Nodes():
        if node.Name() == "PRT3_Random":
            randNodes.append(node)
    return randNodes

def addRandomNode(actor):
    material = actor.Materials()[0]
    st = material.ShaderTree()
    node = st.CreateNode(poser.kNodeTypeCodeSIMPLECOLOR)
    node.SetName("PRT3_Random")
    node.SetInputsCollapsed(1)
    node.SetLocation(160,60)
    import random
    node.Input(0).SetColor( random.random(), random.random(), random.random() )
    return node

def updateRandomNodes(actor, copy=0):
    for node in findRandomNodes(actor):
        kind = actor.Parameter(PARM_TYPE).Value()
        life = 0.0
        if kind == PARTICLE_TYPE:
            life = actor.Parameter(PARM_LIFE).Value()
        if life <= 0 or copy:
            import random
            node.Input(0).SetColor( random.random(), random.random(), random.random() )
        else:
            r, g, b = node.Input(0).Value()
            node.Input(0).SetColor(r, g, b)

def addColorBy(actor, node):
    st = actor.Materials()[0].ShaderTree()
    ramp = st.CreateNode(poser.kNodeTypeCodeCOLORRAMP)
    ramp.SetLocation(160,60)
    node.ConnectToInput(ramp.InputByInternalName("Input"))

def addColorByLife(actor):
    lifeNode = findLifeNode(actor)
    if not lifeNode:
        parm = actor.Parameter(PARM_TYPE)
        lifeNode = addLifeNode(actor, parm.Value())
    addColorBy(actor, lifeNode)

def addColorByRandom(actor):
    randNode = addRandomNode(actor)
    addColorBy(actor, randNode)

def updatePrt3Nodes(actor, copy=0):
    updateLifeNode(actor)
    updateRandomNodes(actor, copy=copy)

def setToAnimate(actor):
    material = actor.Materials()[0]
    st = material.ShaderTree()
    randomNodes = findRandomNodes(actor)
    for randomNode in randomNodes:
        randomNode.Input(0).SetAnimated(1)
    lifeNode = findLifeNode(actor)
    if lifeNode:
        lifeNode.InputByInternalName("Value_1").SetAnimated(1)

def simulate(figure):
    for bodypart in figure.Actors():
        setToAnimate(bodypart)
    scene = poser.Scene()
    for frame in range(scene.NumFrames()):
        scene.SetFrame(int(frame))
        for bodypart in figure.Actors():
            updatePrt3Nodes(bodypart)

def debug():
    pass
